


module udp_rx#(
    parameter                   SOURCE_PORT = 16'd5001      
)(
    input                       i_clk                       ,
    input                       i_rst                       ,
    input                       i_set_source_port_valid     ,
    input  [15: 0]              i_set_source_port           ,
    input                       s_axis_tvalid               ,
    input  [7 : 0]              s_axis_tdata                ,
    input                       s_axis_tlast                ,
    input  [71: 0]              s_axis_tuser                ,
    output                      m_axis_tvalid               ,
    output [7 : 0]              m_axis_tdata                ,
    output                      m_axis_tlast                ,
    output [63: 0]              m_axis_tuser                  
);

    localparam  UDP_PKG_TYPE     = 8'd17                    ; 

    logic   [15: 0]             set_source_port             ;
    logic   [7 : 0]             udp_axis_rx_data            ;
    logic                       udp_axis_rx_last            ;
    logic                       udp_axis_rx_valid           ;
    logic   [15: 0]             rx_cnt                      ;
    logic   [15: 0]             udp_pkg_len                 ;
    logic   [15: 0]             udp_data_len                ;
    logic   [31: 0]             udp_rx_source_ip            ;
    logic   [15: 0]             udp_rx_target_port          ;
    logic   [15: 0]             udp_rx_source_port          ;
    logic   [7 : 0]             axis_rx_data                ;
    logic   [63: 0]             axis_rx_user                ;
    logic                       axis_rx_last                ;
    logic                       axis_rx_valid               ;

    // ila128_bit ila128_bit_udp_rx(
    //     .clk    ( i_clk),
    //     .probe0 ( {s_axis_tvalid,s_axis_tdata,s_axis_tlast,rx_cnt,udp_rx_target_port,udp_rx_source_port,
    //     udp_data_len,ip_pkg_len,m_axis_tvalid,m_axis_tlast,m_axis_tdata})

    // );

    assign m_axis_tdata = axis_rx_data; 
    assign m_axis_tuser = axis_rx_user; 
    assign m_axis_tlast = axis_rx_last; 
    assign m_axis_tvalid = axis_rx_valid;  

    always_ff @(posedge i_clk) begin
        if(i_rst) begin
            set_source_port  <= SOURCE_PORT;
        end
        else if(i_set_source_port_valid) begin
            set_source_port  <= i_set_source_port ;
        end
        else begin
            set_source_port  <= set_source_port;
        end
    end

    always_ff @(posedge i_clk) begin
        if(i_rst) begin
            udp_axis_rx_data  <= 0;
            udp_axis_rx_last  <= 0;
            udp_axis_rx_valid <= 0;
        end 
        else begin
            udp_axis_rx_data  <= s_axis_tdata ;
            udp_axis_rx_last  <= s_axis_tlast ;
            udp_axis_rx_valid <= s_axis_tvalid;
        end
    end

    always_ff @(posedge i_clk) begin
        if(i_rst) 
            udp_pkg_len <= 0;
        else if(s_axis_tvalid) 
            udp_pkg_len <= s_axis_tuser[15:0];
        else   
            udp_pkg_len <= udp_pkg_len;
    end

    always_ff @(posedge i_clk) begin
        if(i_rst) 
            udp_data_len <= 0;
        else if(s_axis_tvalid) 
            udp_data_len <= s_axis_tuser[15:0]-8;
        else   
            udp_data_len <= udp_data_len;
    end
    
    always_ff @(posedge i_clk) begin
        if(i_rst)
            udp_rx_source_ip <= 0;
        else if(s_axis_tvalid)
            udp_rx_source_ip <= s_axis_tuser[55:24];
        else 
            udp_rx_source_ip <= udp_rx_source_ip;
    end 

    always_ff @(posedge i_clk) begin
        if(i_rst)
            rx_cnt <= 0;
        else if(udp_axis_rx_valid)
            rx_cnt <= rx_cnt + 1;
        else 
            rx_cnt <= 0;
    end

    always_ff @(posedge i_clk) begin
        if(i_rst)
            udp_rx_source_port <= 0;
        else if(rx_cnt >= 0 && rx_cnt <= 1 && udp_axis_rx_valid == 1)
            udp_rx_source_port <= {udp_rx_source_port[7:0],udp_axis_rx_data};
        else 
            udp_rx_source_port <= udp_rx_source_port;
    end 

    always_ff @(posedge i_clk) begin
        if(i_rst)
            udp_rx_target_port <= 0;
        else if(rx_cnt >= 2 && rx_cnt <= 3 && udp_axis_rx_valid == 1)
            udp_rx_target_port <= {udp_rx_target_port[7:0],udp_axis_rx_data};
        else 
            udp_rx_target_port <= udp_rx_target_port;
    end 
 
    always_ff @(posedge i_clk) begin
        if(i_rst)
            axis_rx_data <= 0;
        else 
            axis_rx_data <= udp_axis_rx_data;
    end

    always_ff @(posedge i_clk) begin
        if(i_rst)
            axis_rx_user <= 'd0;
        else 
            axis_rx_user <= {udp_rx_source_port, udp_rx_source_ip, udp_data_len};
    end
 
    always_ff @(posedge i_clk) begin
        if(i_rst)
            axis_rx_last <= 'd0;
        else if(udp_axis_rx_last & udp_axis_rx_valid)
            axis_rx_last <= 'd1;
        else 
            axis_rx_last <= 'd0;
    end

    always_ff @(posedge i_clk) begin
        if(i_rst)
            axis_rx_valid <= 0;
        else if(axis_rx_last && axis_rx_valid)
            axis_rx_valid <= 0;
        else if(rx_cnt == 8 && udp_rx_target_port == set_source_port && s_axis_tuser[23:16] == UDP_PKG_TYPE)
            axis_rx_valid <= 1;
        else 
            axis_rx_valid <= axis_rx_valid;
    end



endmodule
